"""
事件分类模块（SVM模型）
功能：基于支持向量机对事件进行多分类
分类类型：拥堵、施工、交通事故、设施故障、其他
输入：预处理后的文本特征
输出：事件分类标签
"""

import numpy as np
from typing import List, Dict, Tuple
from sklearn.svm import SVC
from sklearn.feature_extraction.text import TfidfVectorizer

class EventClassifier:
    def __init__(self):
        # 分类标签
        self.labels = ['拥堵', '施工', '交通事故', '设施故障', '其他']
        # SVM分类器
        self.model = SVC(kernel='rbf', probability=True)
        # TF-IDF向量化器
        self.vectorizer = TfidfVectorizer()
        # 是否已训练
        self.is_trained = False
    
    def train(self, texts: List[str], labels: List[str]) -> bool:
        """
        训练分类模型
        :param texts: 训练文本列表
        :param labels: 对应标签列表
        :return: 是否训练成功
        """
        try:
            # 验证标签
            for label in labels:
                if label not in self.labels:
                    raise ValueError(f"无效标签: {label}")
            
            # 文本向量化
            X = self.vectorizer.fit_transform(texts)
            
            # 训练模型
            self.model.fit(X, labels)
            self.is_trained = True
            return True
        except Exception as e:
            print(f"模型训练失败: {e}")
            return False
    
    def predict(self, text: str) -> Tuple[str, float]:
        """
        预测单个文本的分类
        :param text: 待分类文本
        :return: (分类标签, 置信度)
        """
        if not self.is_trained:
            raise RuntimeError("模型未训练，请先调用train方法")
            
        # 文本向量化
        X = self.vectorizer.transform([text])
        
        # 预测
        prediction = self.model.predict(X)[0]
        probabilities = self.model.predict_proba(X)[0]
        confidence = max(probabilities)
        
        return prediction, confidence
    
    def predict_batch(self, texts: List[str]) -> List[Tuple[str, float]]:
        """
        批量预测文本分类
        :param texts: 待分类文本列表
        :return: [(分类标签, 置信度), ...]
        """
        if not self.is_trained:
            raise RuntimeError("模型未训练，请先调用train方法")
            
        # 文本向量化
        X = self.vectorizer.transform(texts)
        
        # 预测
        predictions = self.model.predict(X)
        probabilities = self.model.predict_proba(X)
        
        results = []
        for i, pred in enumerate(predictions):
            confidence = max(probabilities[i])
            results.append((pred, confidence))
            
        return results
    
    def get_labels(self) -> List[str]:
        """
        获取分类标签列表
        :return: 分类标签列表
        """
        return self.labels

# 使用示例
def main():
    classifier = EventClassifier()
    
    # 训练数据
    train_texts = [
        "道路拥堵严重，车辆行驶缓慢",
        "南京路进行道路施工，预计工期一个月",
        "在中山路与解放路交叉口发生一起两车相撞事故",
        "地铁2号线人民广场站电梯故障停运",
        "今天天气很好，适合出行",
        "市中心交通管制，绕行路线已发布",
        "淮海路路面塌陷，紧急抢修中",
        "公交线路调整，新增站点",
        "共享单车投放量增加，缓解出行压力",
        "雾霾天气，能见度低"
    ]
    
    train_labels = [
        '拥堵', '施工', '交通事故', '设施故障', '其他',
        '拥堵', '施工', '其他', '其他', '其他'
    ]
    
    # 训练模型
    if classifier.train(train_texts, train_labels):
        print("模型训练成功")
        
        # 测试数据
        test_texts = [
            "解放路发生严重交通事故，多车连环相撞",
            "地铁1号线因设备故障临时停运",
            "高架桥车流量大，通行缓慢"
        ]
        
        # 批量预测
        results = classifier.predict_batch(test_texts)
        
        print("\n分类结果:")
        for i, (label, confidence) in enumerate(results):
            print(f"  文本 {i+1}: {label} (置信度: {confidence:.2f})")
            print(f"    内容: {test_texts[i]}")
    else:
        print("模型训练失败")

if __name__ == "__main__":
    main()